import os
import dotenv

import requests
from langchain.agents import initialize_agent, Tool, AgentType
from langchain_openai import ChatOpenAI
from langchain.tools import tool

dotenv.load_dotenv()

# 定义Bocha Web Search工具
@tool
def bocha_web_search_tool(query: str, count: int = 8) -> str:
    """
    使用Bocha Web Search API进行联网搜索，返回搜索结果的字符串。

    参数:
    - query: 搜索关键词
    - count: 返回的搜索结果数量

    返回:
    - 搜索结果的字符串形式
    """
    url = 'https://api.bochaai.com/v1/web-search'
    headers = {
        'Authorization': f'Bearer {os.getenv("BOCHA_API_KEY")}',  # 请替换为你的API密钥
        'Content-Type': 'application/json'
    }
    data = {
        "query": query,
        "freshness": "noLimit",  # 搜索的时间范围，例如 "oneDay", "oneWeek", "oneMonth", "oneYear", "noLimit"
        "summary": True,  # 是否返回长文本摘要总结
        "count": count
    }

    response = requests.post(url, headers=headers, json=data)

    if response.status_code == 200:
        # 返回给大模型的格式化的搜索结果文本
        # 可以自己对博查的搜索结果进行自定义处理
        return str(response.json())
    else:
        raise Exception(f"API请求失败，状态码: {response.status_code}, 错误信息: {response.text}")


# 创建LangChain工具
bocha_tool = Tool(
    name="BochaWebSearch",
    func=bocha_web_search_tool,
    description="使用Bocha Web Search API进行网络搜索"
)

# 初始化OpenAI语言模型
llm = ChatOpenAI(
    model="kimi-k2-0711-preview",
    temperature=0
)

# 初始化代理，设置Agent自动决定何时调用工具
tools = [bocha_tool]  # 包含你定义的Bocha Web Search工具

agent = initialize_agent(
    tools=tools,
    agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
    llm=llm,
    verbose=False
)

# 测试用例
query = "阿里巴巴2024年的ESG报告主要讲了哪些内容？"
response = agent.run(f"请帮我搜索以下内容：{query}")
print(response)
